# coding = utf-8

'''
模型test相关介绍
'''

from network.HDenseUnet import dense_rnn_net
from dataset.IRCAD_HDENSE_UNET import IRCAD_DATA
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
import torch
from tqdm import tqdm
import numpy as np
import json
import os
from scipy import ndimage
from torch.nn.functional import softmax
from skimage import measure
from pathlib2 import Path

def load_liver_mask(roi_file):
    with open(roi_file, "r") as file:
         roi = json.load(file)
    return roi

def predict_tumor_inwindow(model, imgs_test, num, mini, maxi, batch, input_size, input_cols):

    batch = batch
    img_deps = input_size
    img_rows = input_size
    img_cols = input_cols


    window_cols = int((img_cols/4))
    count = 0
    box_test = np.zeros((batch,img_deps,img_rows,img_cols, 1), dtype="float32")

    x = imgs_test.shape[0]
    y = imgs_test.shape[1]
    z = imgs_test.shape[2]
    right_cols = int(min(z,maxi[2]+10)-img_cols)
    left_cols  = max(0,min(mini[2]-5, right_cols))
    score = np.zeros((x, y, z, num), dtype= 'float32')
    score_num = np.zeros((x, y, z, num), dtype= 'int16')

    for cols in range(left_cols,right_cols+window_cols,window_cols):
        if cols > z - img_cols:
            patch_test = imgs_test[0:img_deps, 0:img_rows, z-img_cols:z]
            box_test[count, :, :, :, 0] = patch_test
            # print ('final', img_cols-window_cols, img_cols)
            patch_test_mask = model.predict(box_test, batch_size=batch, verbose=0)
            patch_test_mask = K.softmax(patch_test_mask)
            patch_test_mask = K.eval(patch_test_mask)
            patch_test_mask = patch_test_mask[:,:,:,1:-1,:]

            for i in xrange(batch):
                score[0:img_deps, 0:img_rows,  z-img_cols+1:z-1, :] += patch_test_mask[i]
                score_num[0:img_deps, 0:img_rows,  z-img_cols+1:z-1, :] += 1
        else:
            patch_test = imgs_test[0:img_deps, 0:img_rows, cols:cols + img_cols]
            box_test[count, :, :, :, 0] = patch_test
            patch_test_mask = model.predict(box_test, batch_size=batch, verbose=0)
            patch_test_mask = K.softmax(patch_test_mask)
            patch_test_mask = K.eval(patch_test_mask)
            patch_test_mask = patch_test_mask[:,:,:,1:-1,:]
            for i in xrange(batch):
                score[0:img_deps, 0:img_rows, cols+1:cols+img_cols-1, :] += patch_test_mask[i]
                score_num[0:img_deps, 0:img_rows, cols+1:cols+img_cols-1, :] += 1
    score = score/(score_num+1e-4)
    score1 = score[:,:,:,num-2]
    score2 = score[:,:,:,num-1]
    return score1, score2

def predict_tumor_inwindowv2(model, imgs_test, num, mini, maxi, batch, input_size, input_cols):

    batch = batch
    img_deps = input_size
    img_rows = input_size
    img_cols = input_cols


    window_cols = int((img_cols/4))
    count = 0
    box_test = np.zeros((batch,img_deps,img_rows,img_cols, 1), dtype="float32")

    x = imgs_test.shape[0]
    y = imgs_test.shape[1]
    z = imgs_test.shape[2]
    right_cols = int(min(z,maxi[2]+10)-img_cols)
    left_cols  = max(0,min(mini[2]-5, right_cols))
    score = np.zeros((x, y, z, num), dtype= 'float32')
    score_num = np.zeros((x, y, z, num), dtype= 'int16')

    for cols in range(left_cols,right_cols+window_cols,window_cols):
        if cols > z - img_cols:
            patch_test = imgs_test[0:img_deps, 0:img_rows, z-img_cols:z]
            #box_test[count, :, :, :, 0] = patch_test
            #print(box_test.shape)

            patch_test = patch_test.transpose([2, 0, 1])
            patch_test = torch.from_numpy(patch_test).type(torch.FloatTensor)
            patch_test = patch_test.unsqueeze(0).cuda()
            with torch.no_grad():
                patch_test_output = model(patch_test)
            patch_test_output = softmax(patch_test_output, dim=1)
            patch_test_output = patch_test_output[:, :, 1:-1, :, :]
            patch_test_output = patch_test_output.cpu().numpy()
            patch_test_output = patch_test_output.transpose([0, 3, 4, 2, 1])

            #print(patch_test_output.shape, score[0:img_deps, 0:img_rows,  z-img_cols+1:z-1, :].shape)

            # print ('final', img_cols-window_cols, img_cols)
            #patch_test_mask = model.predict(box_test, batch_size=batch, verbose=0)
            #patch_test_mask = K.softmax(patch_test_mask)
            #patch_test_mask = K.eval(patch_test_mask)
            #patch_test_mask = patch_test_mask[:,:,:,1:-1,:]

            for i in range(batch):
                score[0:img_deps, 0:img_rows,  z-img_cols+1:z-1, :] += patch_test_output[i]
                score_num[0:img_deps, 0:img_rows,  z-img_cols+1:z-1, :] += 1
        else:
            patch_test = imgs_test[0:img_deps, 0:img_rows, cols:cols + img_cols]
            patch_test = patch_test.transpose([2,0,1])
            patch_test = torch.from_numpy(patch_test).type(torch.FloatTensor)
            patch_test = patch_test.unsqueeze(0).cuda()

            with torch.no_grad():
                patch_test_output = model(patch_test)
            patch_test_output = softmax(patch_test_output, dim=1)

            patch_test_output = patch_test_output[:,:,1:-1,:,:]
            patch_test_output = patch_test_output.cpu().numpy()
            patch_test_output = patch_test_output.transpose([0,3,4,2,1])


            #box_test[count, :, :, :, 0] = patch_test
            #patch_test_mask = model.predict(box_test, batch_size=batch, verbose=0)
            #patch_test_mask = K.softmax(patch_test_mask)
            #patch_test_mask = K.eval(patch_test_mask)
            #patch_test_mask = patch_test_mask[:,:,:,1:-1,:]
            for i in range(batch):
                score[0:img_deps, 0:img_rows, cols+1:cols+img_cols-1, :] += patch_test_output[i]
                score_num[0:img_deps, 0:img_rows, cols+1:cols+img_cols-1, :] += 1

    score = score/(score_num+1e-4)
    score1 = score[:,:,:,num-2]
    score2 = score[:,:,:,num-1]
    return score1, score2

def test():
    model_file = "/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/fusion/epoll_199.pkl"
    model = torch.load(model_file).cuda()
    roi_path = "/home/diaozhaoshuo/log/BeliefFunctionNN/3diradb_chengkung/output"
    save_path = "/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/predict"

    for j in range(20):
        data_path = "/datasets/3Dircadb/origion/case_{}".format(str(j).zfill(5))
        origion_data_path = os.path.join(data_path, "image.npy")
        mask_path = os.path.join(data_path, "segment.npy")
        image = np.load(origion_data_path)
        mask = np.load(mask_path)


        roi_file = os.path.join(roi_path, "predict_{}.npy".format(str(j).zfill(5)))
        roi = np.load(roi_file)
        roi = ndimage.binary_dilation(roi, iterations=1).astype(roi.dtype)

        image = image.transpose([1,2,0])
        mask = mask.transpose([1,2,0])
        roi = roi.transpose([1,2,0])



        index = np.where(roi == 1)
        mini = np.min(index, axis=-1)
        maxi = np.max(index, axis=-1)

        score1, score2 = predict_tumor_inwindowv2(model=model, imgs_test=image, num=3, mini=mini, maxi=maxi, batch=1, input_size=512,
                               input_cols=8)

        case_path = Path(save_path) / "case_{}".format(str(j).zfill(5))
        if not case_path.exists():
            case_path.mkdir(parents=True)

        np.save(os.path.join(save_path,"case_{}/score1.npy".format(str(j).zfill(5))),score1)
        np.save(os.path.join(save_path,"case_{}/score2.npy".format(str(j).zfill(5))),score2)


        result1 = score1
        result2 = score2
        result1[result1 >= 0.5] = 1
        result1[result1 < 0.5] = 0
        result2[result2 >= 0.9] = 1
        result2[result2 < 0.9] = 0
        result1[result2 == 1] = 1

        mask1 = mask.copy()
        mask2 = mask.copy()
        mask1[mask == 2] = 1
        mask2[mask == 1] = 0

        dice1 = (result1*mask1).sum() / (result1.sum() + mask1.sum() + 1e-4)
        dice2 = (result2*mask2).sum() / (mask2.sum() + result2.sum() + 1e-4)

        print(j, 2*dice1, 2*dice2)

        Segmask = result2
        box = []
        [liver_res, num] = measure.label(result1, return_num=True)
        region = measure.regionprops(liver_res)
        for i in range(num):
            box.append(region[i].area)
        label_num = box.index(max(box)) + 1
        liver_res[liver_res != label_num] = 0
        liver_res[liver_res == label_num] = 1

        #  preserve the largest liver
        mask = ndimage.binary_dilation(mask, iterations=1).astype(mask.dtype)
        box = []
        [liver_labels, num] = measure.label(mask, return_num=True)
        region = measure.regionprops(liver_labels)
        for i in range(num):
            box.append(region[i].area)
        label_num = box.index(max(box)) + 1
        liver_labels[liver_labels != label_num] = 0
        liver_labels[liver_labels == label_num] = 1
        liver_labels = ndimage.binary_fill_holes(liver_labels).astype(int)

        #  preserve tumor within ' largest liver' only
        Segmask = Segmask * liver_labels
        Segmask = ndimage.binary_fill_holes(Segmask).astype(int)
        Segmask = np.array(Segmask, dtype='uint8')
        liver_res = np.array(liver_res, dtype='uint8')
        liver_res = ndimage.binary_fill_holes(liver_res).astype(int)
        liver_res[Segmask == 1] = 2
        liver_res = np.array(liver_res, dtype='uint8')

        dice1 = (result1 * mask1).sum()  / (result1.sum() + mask1.sum() + 1e-4)
        dice2 = (Segmask * mask2).sum()  / (mask2.sum() + Segmask.sum() +1e-4)
        print(j, 2 * dice1, 2 * dice2)




def testv2():
    model_file = "/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/fusion/epoll_199.pkl"
    model = torch.load(model_file).cuda()
    roi_path = "/home/diaozhaoshuo/log/BeliefFunctionNN/3diradb_chengkung/output"

    for i in range(20):
        data_path = "/datasets/3Dircadb/origion/case_{}".format(str(i).zfill(5))
        origion_data_path = os.path.join(data_path, "image.npy")
        mask_path = os.path.join(data_path, "segment.npy")
        image = np.load(origion_data_path)
        mask = np.load(mask_path)

        roi_file = os.path.join(roi_path, "predict_{}.npy".format(str(i).zfill(5)))
        roi = np.load(roi_file)
        roi = ndimage.binary_dilation(roi, iterations=1).astype(roi.dtype)

        image = image.transpose([1, 2, 0])
        mask = mask.transpose([1, 2, 0])
        roi = roi.transpose([1, 2, 0])
        result1 = np.load("score1.npy")
        result2 = np.load("score2.npy")
        result1[result1 >= 0.5] = 1
        result1[result1 < 0.5] = 0
        result2[result2 >= 0.9] = 1
        result2[result2 < 0.9] = 0
        result1[result2 == 1] = 1

        Segmask = result2
        box = []
        [liver_res, num] = measure.label(result1, return_num=True)
        region = measure.regionprops(liver_res)
        for i in range(num):
            box.append(region[i].area)
        label_num = box.index(max(box)) + 1
        liver_res[liver_res != label_num] = 0
        liver_res[liver_res == label_num] = 1

        #  preserve the largest liver
        mask = ndimage.binary_dilation(mask, iterations=1).astype(mask.dtype)
        box = []
        [liver_labels, num] = measure.label(mask, return_num=True)
        region = measure.regionprops(liver_labels)
        for i in range(num):
            box.append(region[i].area)
        label_num = box.index(max(box)) + 1
        liver_labels[liver_labels != label_num] = 0
        liver_labels[liver_labels == label_num] = 1
        liver_labels = ndimage.binary_fill_holes(liver_labels).astype(int)

        #  preserve tumor within ' largest liver' only
        Segmask = Segmask * liver_labels
        Segmask = ndimage.binary_fill_holes(Segmask).astype(int)
        Segmask = np.array(Segmask, dtype='uint8')
        liver_res = np.array(liver_res, dtype='uint8')
        liver_res = ndimage.binary_fill_holes(liver_res).astype(int)
        liver_res[Segmask == 1] = 2
        liver_res = np.array(liver_res, dtype='uint8')

        print(np.unique(liver_res), np.unique(Segmask))


        mask1 = mask.copy()
        mask2 = mask.copy()
        mask1[mask == 2] = 1
        mask2[mask == 1] = 0
        mask2[mask == 2] = 1

        dice1 = (result1 * mask1).sum() / (result1.sum() + mask1.sum())
        dice2 = (Segmask * mask2).sum() / (mask2.sum() + Segmask.sum())

        print(2 * dice1, 2 * dice2)

        break


if __name__ == '__main__':
    test()